from docplex.mp.model import Model
from QKDNetwork import QKDNetwork
import time


def calculateFlowPath(net: QKDNetwork) -> list:  # 求解候选路径
    # return [{"sd": [1, 2, 4], "path": [1, 3, 2]}...]
    model = Model(name="cal_flow")
    T = model.continuous_var(lb=0, name="T")

    p_dim1, p_dim2 = net.num_nodes + 1, net.num_nodes + 1
    q_dim1, q_dim2 = len(net.sd_list), net.max_candidate_size

    # 节点u向节点v发射的时隙数量
    p = model.continuous_var_dict(
        [(i, j) for i in range(p_dim1) for j in range(p_dim2)], lb=0, name="p"
    )

    # 请求i是否选择了第j条候选路径作为自己的密钥同步路径
    q = model.binary_var_dict(
        [(i, j) for i in range(q_dim1) for j in range(q_dim2)], lb=0, name="q"
    )

    big_M = 10000

    # 7a ✅
    for sd_index, sd_item in enumerate(net.sd_list):
        i = sd_index
        sole_path = model.linear_expr()
        for candidate_path_index, candidate_path_item in enumerate(
            net.candidate_paths[sd_index]["paths"]
        ):
            j = candidate_path_index
            sole_path += q[i, j]
        model.add_constraint(sole_path == 1)

    def is_sublist(sub, main):
        sub_length = len(sub)
        return any(
            sub == main[i : i + sub_length] for i in range(len(main) - sub_length + 1)
        )

    # 7b ✅
    for edge in net.G.edges:
        u = edge[0]
        v = edge[1]
        flow = model.linear_expr()
        flow += net.G.nodes[u]["transmitter_rate"] * p[u, v]
        flow += net.G.nodes[v]["transmitter_rate"] * p[v, u]
        flow += net.G[u][v]["remain_key_num"]
        for sd_index, sd_item in enumerate(net.sd_list):
            i = sd_index
            for j, path in enumerate(net.candidate_paths[sd_index]["paths"]):
                if is_sublist([u, v], path) or is_sublist([v, u], path):
                    # (u, v)在链路上
                    flow -= (sd_item[2] - sd_item[3]) * q[i, j]
        model.add_constraint(flow >= 0)

    # 7c ✅
    for edge in net.G.edges:
        u, v = edge[0], edge[1]
        S_u, S_v = net.G.nodes[u]["transmitter"], net.G.nodes[v]["transmitter"]
        model.add_constraint(p[u, v] <= big_M * S_u)
        model.add_constraint(p[v, u] <= big_M * S_v)

    # 7d ✅
    for node in net.G.nodes:
        if net.G.nodes[node]["transmitter"]:
            slot_expr = model.linear_expr()
            for node_neighbor in net.G.neighbors(node):
                u, v = node, node_neighbor
                slot_expr += p[u, v]
            model.add_constraint(slot_expr <= T)

    model.minimize(T)
    sol = model.solve()
    ret = []
    p_normalized = {}

    if sol:
        for (i, j), variable in q.items():
            if variable.solution_value >= 0.95:
                ret.append(
                    {"sd": net.sd_list[i], "path": net.candidate_paths[i]["paths"][j]}
                )
        for node in net.G.nodes:
            if net.G.nodes[node]["transmitter"]:
                node_p_values = {}
                total_p = 0
                # 计算到所有邻居的p值总和
                for neighbor in net.G.neighbors(node):
                    p_value = p[node, neighbor].solution_value
                    if p_value > 0:  # 只考虑大于0的值
                        node_p_values[neighbor] = p_value
                        total_p += p_value

                if total_p > 0:
                    normalized_values = {}
                    for neighbor, p_value in node_p_values.items():
                        normalized_values[neighbor] = p_value / total_p
                    p_normalized[node] = normalized_values

    # 需要返回一个预估的剩余时间数

    return ret, p_normalized, T.solution_value


def estimateSlotNum(net: QKDNetwork) -> list:  # 求解候选路径
    # return [{"sd": [1, 2, 4], "path": [1, 3, 2]}...]
    model = Model(name="estimate_slot_num")
    T = model.integer_var(lb=0, name="T")

    p_dim1, p_dim2 = net.num_nodes + 1, net.num_nodes + 1
    q_dim1, q_dim2 = len(net.sd_list), net.max_candidate_size

    # 节点u向节点v发射的时隙数量
    p = model.integer_var_dict(
        [(i, j) for i in range(p_dim1) for j in range(p_dim2)], lb=0, name="p"
    )

    # 请求i是否选择了第j条候选路径作为自己的密钥同步路径
    q = model.binary_var_dict(
        [(i, j) for i in range(q_dim1) for j in range(q_dim2)], lb=0, name="q"
    )

    big_M = 10000

    # 7a ✅
    for sd_index, sd_item in enumerate(net.sd_list):
        i = sd_index
        sole_path = model.linear_expr()
        for candidate_path_index, candidate_path_item in enumerate(
            net.candidate_paths[sd_index]["paths"]
        ):
            j = candidate_path_index
            sole_path += q[i, j]
        model.add_constraint(sole_path == 1)

    def is_sublist(sub, main):
        sub_length = len(sub)
        return any(
            sub == main[i : i + sub_length] for i in range(len(main) - sub_length + 1)
        )

    # 7b ✅
    for edge in net.G.edges:
        u = edge[0]
        v = edge[1]
        flow = model.linear_expr()
        flow += net.G.nodes[u]["transmitter_rate"] * p[u, v]
        flow += net.G.nodes[v]["transmitter_rate"] * p[v, u]
        flow += net.G[u][v]["remain_key_num"]
        for sd_index, sd_item in enumerate(net.sd_list):
            i = sd_index
            for j, path in enumerate(net.candidate_paths[sd_index]["paths"]):
                if is_sublist([u, v], path) or is_sublist([v, u], path):
                    # (u, v)在链路上
                    flow -= (sd_item[2] - sd_item[3]) * q[i, j]
        model.add_constraint(flow >= 0)

    # 7c ✅
    for edge in net.G.edges:
        u, v = edge[0], edge[1]
        S_u, S_v = net.G.nodes[u]["transmitter"], net.G.nodes[v]["transmitter"]
        model.add_constraint(p[u, v] <= big_M * S_u)
        model.add_constraint(p[v, u] <= big_M * S_v)

    # 7d ✅
    for node in net.G.nodes:
        if net.G.nodes[node]["transmitter"]:
            slot_expr = model.linear_expr()
            for node_neighbor in net.G.neighbors(node):
                u, v = node, node_neighbor
                slot_expr += p[u, v]
            model.add_constraint(slot_expr <= T)

    model.minimize(T)
    sol = model.solve()
    ret = []
    p_normalized = {}

    if sol:
        return T.solution_value

    return 0
